import requests
import json
from img_rag import MilvusImageVectorStore  # 修改导入路径
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
from datetime import datetime
import os

# 数据模型定义
class ChatMessage(BaseModel):
    type: str  # 'user' 或 'assistant'
    content: str
    timestamp: str

class ChatRequest(BaseModel):
    message: str
    history: List[ChatMessage] = []

class ChatResponse(BaseModel):
    success: bool
    message: str
    questions: List[str] = []
    images: List[str] = []
    timestamp: str

class ErrorResponse(BaseModel):
    success: bool
    error: str
    timestamp: str


class DeepSeekAPI:
    def __init__(self, api_key: str, base_url: str = "https://api.deepseek.com"):
        self.api_key = api_key
        self.base_url = base_url
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        # 初始化图片向量存储，使用正确的collection_name
        self.image_store = MilvusImageVectorStore(collection_name="img_info")

    def build_prompt(self, user_message: str, chat_history: List[ChatMessage], image_results: List[dict]) -> tuple:
        """构建系统提示词和用户消息"""
        # 构造系统提示
        system_prompt = """你是一位火星科普大模型。请整合用户的提问，生成一篇连贯的文章。
                要求：
                1. 请结合检索到图片和你已有的知识库生成一篇完整的文章，结构清晰，包含引言、主体和结论
                2. 文章语言要通俗易懂，适合科普阅读
                3. 在文章结尾提供三个基于内容的相关提问
                4. 如果提供了相关图片，请在适当位置引用这些图片，只引用用户给的图片链接，图片用 \\image(url)来表示
                6. 在文章结尾单独列出一段话，内容是生成三个相关提问，格式为： ## 相关提问 1. 提问1xxx 2. 提问2xxx 3. 提问3xxx。"""

        # 添加聊天历史
        history_text = ""
        if chat_history:
            history_text += "对话历史：\n"
            for msg in chat_history[-5:]:  # 只取最近5条消息
                role = "用户" if msg.type == "user" else "助手"
                history_text += f"{role}: {msg.content}\n"
            history_text += "\n"

        # 构造用户消息，包含图片信息
        image_info = ""
        if image_results:
            image_info = "\n\n相关图片:\n"
            for i, img in enumerate(image_results):
                image_info += f"![图片{i + 1}]({img['url']})\n"

        user_message = f"{history_text}用户问题: {user_message}{image_info}"

        return system_prompt, user_message

    def parse_response(self, content: str) -> dict:
        """解析AI响应，提取回答和相关问题"""
        questions = []
        message = content

        # 提取相关提问
        if "## 相关提问" in content:
            message = content.split("## 相关提问")[0].strip()
            question_section = content.split("## 相关提问")[-1].strip()
            for line in question_section.splitlines():
                line = line.strip()
                if line.startswith(("1.", "2.", "3.")):
                    questions.append(line[2:].strip())

        return {
            "message": message,
            "questions": questions
        }

    def generate_with_images(self, query: str, top_k: int = 2) -> dict:
        """结合图片信息和DeepSeek知识库生成回答"""
        # 搜索相关图片
        image_results = self.image_store.search_with_description_and_url(query, top_k)

        # 构建提示词
        system_prompt, user_message = self.build_prompt(query, [], image_results)

        # 调用DeepSeek API
        payload = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message}
            ],
            "stream": False
        }

        try:
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=self.headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            result = response.json()
            content = result['choices'][0]['message']['content']

            # 解析响应
            parsed_result = self.parse_response(content)
            # 提取图片URL列表
            image_urls = [img['url'] for img in image_results]

            return {
                "message": parsed_result["message"],
                "questions": parsed_result["questions"],
                "images": image_urls
            }
        except Exception as e:
            return {
                "message": f"调用DeepSeek API时出错: {str(e)}",
                "questions": [],
                "images": []
            }


# 创建FastAPI应用
app = FastAPI(
    title="AI聊天助手API - DeepSeek版本",
    description="火星地图服务的AI聊天助手后端API（集成DeepSeek和图片生成）",
    version="1.0.0"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境中应该指定具体的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化DeepSeek API客户端
# 请确保设置环境变量 DEEPSEEK_API_KEY 或在此处直接提供API密钥
DEEPSEEK_API_KEY = "sk-97a3ba628ed54f0dab2d10bdddd1d71f"
deepseek_client = DeepSeekAPI(DEEPSEEK_API_KEY)


# 聊天API端点
@app.post("/api/chat2", response_model=ChatResponse)
async def chat2_endpoint(request: ChatRequest):
    """
    处理聊天请求的API端点（DeepSeek版本，支持图片生成和RAG）
    """
    try:

        image_results = deepseek_client.image_store.search_with_description_and_url(
            request.message, top_k=3
        )
        print(f"检索到的图片结果: {image_results}")  # 添加这行来查看实际返回内容

        # 构建提示词
        system_prompt, user_message = deepseek_client.build_prompt(
            request.message, request.history, image_results
        )

        # 调用DeepSeek API
        payload = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message}
            ],
            "stream": False
        }

        response = requests.post(
            f"{deepseek_client.base_url}/chat/completions",
            headers=deepseek_client.headers,
            json=payload,
            timeout=60
        )
        response.raise_for_status()
        result = response.json()
        content = result['choices'][0]['message']['content']

        # 解析响应
        parsed_result = deepseek_client.parse_response(content)
        # 提取图片URL列表
        image_urls = [img['url'] for img in image_results]

        # 返回响应
        return ChatResponse(
            success=True,
            message=parsed_result["message"],
            questions=parsed_result["questions"],
            images=image_urls,
            timestamp=datetime.now().isoformat(),
        )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail={
                "success": False,
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
        )


# 健康检查端点
@app.get("/health")
async def health_check():
    """健康检查端点"""
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}


# 根路径
@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "AI聊天助手API - DeepSeek版本",
        "version": "1.0.0",
        "endpoints": {
            "chat": "/api/chat2",
            "health": "/health",
            "docs": "/docs"
        }
    }


# 启动服务器的代码（如果需要）
if __name__ == "__main__":
    import uvicorn

    print("启动AI聊天助手API服务器（DeepSeek版本）...")
    print("访问地址: http://localhost:8001")
    print("API文档: http://localhost:8001/docs")
    print("API端点: POST /api/chat2")
    print("健康检查: GET /health")

    # 启动服务器
    uvicorn.run(
        "popular_fastapi:app",
        host=os.getenv("HOST", "0.0.0.0"),
        port=int(os.getenv("PORT", 8001)),
        reload=True,  # 开发模式下自动重载
        log_level=os.getenv("LOG_LEVEL", "info")
    )

